﻿//-----------------------------------------------------------------------------------------------------------
// Wbfs Engine (WBFSe) Project by Omega Frost 
// http://wbfse.codeplex.com/
//
// WBFSe is Licensed under the terms of the 
// Microsoft Reciprocal License (Ms-RL)
//-----------------------------------------------------------------------------------------------------------
using System;
using System.IO;
using System.Runtime.InteropServices;
using WBFSe.IO;
using System.Text;

namespace WBFSe.OS
{
    /// <summary>
    /// Operating System interop methods
    /// </summary>
    public static class Native
    {
        // Constants

        public static readonly IntPtr InvalidHandle = new IntPtr(-1);

        // Methods

        //------------------------------------------------------------------------------------------
        // Locks a volume
        //------------------------------------------------------------------------------------------
        internal static bool LockVolume(IntPtr handle)
        {
            uint dummy = 0;
            return DeviceIoControl(handle, IOCtrl.FsctlLockVolume, IntPtr.Zero, 0,
                IntPtr.Zero, 0, ref dummy, IntPtr.Zero) != 0;
        }

        //------------------------------------------------------------------------------------------
        // Unlocks a volume
        //------------------------------------------------------------------------------------------
        internal static bool UnlockVolume(IntPtr handle)
        {
            uint dummy = 0;
            return DeviceIoControl(handle, IOCtrl.FsctlUnlockVolume, IntPtr.Zero, 0,
                IntPtr.Zero, 0, ref dummy, IntPtr.Zero) != 0;
        }

        //------------------------------------------------------------------------------------------
        // Dismounts a volume
        //------------------------------------------------------------------------------------------
        internal static bool DismountVolume(IntPtr handle)
        {
            uint dummy = 0;
            return DeviceIoControl(handle, IOCtrl.FsctlDismountVolume, IntPtr.Zero, 0,
                IntPtr.Zero, 0, ref dummy, IntPtr.Zero) != 0;
        }

        //------------------------------------------------------------------------------------------
        // Gets the physical disk's sector size
        //------------------------------------------------------------------------------------------
        internal static bool GetSectorSize(IntPtr handle, out int sectorSize)
        {
            uint bytes = 0;
            DiskGeometry dg = new DiskGeometry();
            IntPtr pdg = Marshal.AllocHGlobal(Marshal.SizeOf(dg));

            int r = DeviceIoControl(handle, IOCtrl.DiskGetDriveGeometry,
                IntPtr.Zero, 0, pdg, (uint)Marshal.SizeOf(dg), ref bytes, IntPtr.Zero);

            dg = (DiskGeometry)Marshal.PtrToStructure(pdg, typeof(DiskGeometry));
            Marshal.FreeHGlobal(pdg);

            sectorSize = dg.BytesPerSector;
            return r != 0;
        }

        //------------------------------------------------------------------------------------------
        // Gets the partition size
        //------------------------------------------------------------------------------------------
        internal static bool GetPartitionSize(IntPtr handle, out long partitionSize)
        {
            uint bytes = 0;
            PartitionInfo pi = new PartitionInfo();
            IntPtr ppi = Marshal.AllocHGlobal(Marshal.SizeOf(pi));

            int r = DeviceIoControl(handle, IOCtrl.DiskGetPartitionInfo,
                IntPtr.Zero, 0, ppi, (uint)Marshal.SizeOf(pi), ref bytes, IntPtr.Zero);

            pi = (PartitionInfo)Marshal.PtrToStructure(ppi, typeof(PartitionInfo));
            Marshal.FreeHGlobal(ppi);

            partitionSize = pi.PartitionLength;
            return r != 0;
        }

        //------------------------------------------------------------------------------------------
        // Marks a region of the file as empty or fills it with zeroes
        //------------------------------------------------------------------------------------------
        internal static bool SetSparseRegion(IntPtr hFile, long position, long count)
        {
            uint temp = 0;
            Region zdi = new Region();

            zdi.FileOffset = position;
            zdi.BeyondFinalZero = position + count;

            IntPtr pzdi = Marshal.AllocHGlobal((int)Region.Size);
            Marshal.StructureToPtr(zdi, pzdi, true);

            bool r = (DeviceIoControl(hFile,
                IOCtrl.FsctlSetZeroData, pzdi, Region.Size,
                IntPtr.Zero, 0, ref temp, IntPtr.Zero) != 0);

            Marshal.FreeHGlobal(pzdi);
            return r;
        }

        //------------------------------------------------------------------------------------------
        // Marks a file as sparse
        //------------------------------------------------------------------------------------------
        internal static bool SetSparseFile(IntPtr hFile)
        {
            uint dummy = 0;
            return DeviceIoControl(hFile, IOCtrl.FsctlSetSparse, IntPtr.Zero, 
                0, IntPtr.Zero, 0, ref dummy, IntPtr.Zero) != 0;
        }

        //------------------------------------------------------------------------------------------
        // Retrieves the allocated disk space of the file
        //------------------------------------------------------------------------------------------
        internal static bool GetSparseSize(IntPtr hFile, out long size)
        {
            size = 0;
            Region range;

            IntPtr ranges = Marshal.AllocHGlobal((int)Region.Size * 1024);
            IntPtr prange = Marshal.AllocHGlobal((int)Region.Size);

            int r;
            uint retsz = 0;

            range.FileOffset = 0;
            if (!GetFileSizeEx(hFile, out range.BeyondFinalZero))
                return false;

            do
            {
                // Send the range to the unmanaged memory
                Marshal.StructureToPtr(range, prange, true);

                // Retrieve the layout
                r = DeviceIoControl(hFile, IOCtrl.FsctlQueryAllocatedRanges,
                    prange, Region.Size, ranges, 1024 * Region.Size,
                    ref retsz, IntPtr.Zero);

                if (r == 0)
                {
                    Marshal.FreeHGlobal(ranges);
                    Marshal.FreeHGlobal(prange);

                    return false;
                }

                Region region;
                IntPtr p = ranges;
                uint n = retsz / Region.Size;

                region.FileOffset = 0;
                region.BeyondFinalZero = 0;

                for (uint i = 0; i < n; i++)
                {
                    region = (Region)Marshal.
                        PtrToStructure(p, typeof(Region));

                    size += region.BeyondFinalZero; //Length

                    p = IntPtr.Add(p, (int)Region.Size);
                }

                long delta = region.FileOffset +
                    region.BeyondFinalZero -
                    range.FileOffset;

                range.FileOffset += delta;
                range.BeyondFinalZero -= delta; //Length

            } while (r == 234); // ERROR_MORE_DATA

            // Free the unmanaged memory
            Marshal.FreeHGlobal(ranges);
            Marshal.FreeHGlobal(prange);

            return true;
        }

        //------------------------------------------------------------------------------------------
        // Retrieves information about a volume
        //------------------------------------------------------------------------------------------
        [DllImport("Kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        internal extern static bool GetVolumeInformation(string RootPathName,
            StringBuilder VolumeNameBuffer, int VolumeNameSize,
            out uint VolumeSerialNumber, out uint MaximumComponentLength,
            out FSFeature FileSystemFlags, StringBuilder FileSystemNameBuffer,
            int nFileSystemNameSize);


        //------------------------------------------------------------------------------------------
        // Retrieves the file size
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll")]
        internal static extern bool GetFileSizeEx([In]IntPtr hFile, [Out]out long lpFileSize);


        //------------------------------------------------------------------------------------------
        // Sets the position of the file pointer
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll")]
        internal static extern bool SetFilePointerEx([In]IntPtr hFile, [In]long liDistanceToMove,
           [Out]out long lpNewFilePointer, [In]SeekOrigin dwMoveMethod);


        //------------------------------------------------------------------------------------------
        // Creates a file handle
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        internal static extern IntPtr CreateFile([In]String lpFileName, 
            [In][MarshalAs(UnmanagedType.U4)]FileAccess dwDesiredAccess,
            [In][MarshalAs(UnmanagedType.U4)]FileShare dwShareMode,
            [In]IntPtr lpSecurityAttributes,
            [In][MarshalAs(UnmanagedType.U4)]FileMode dwCreationDisposition,
            [In][MarshalAs(UnmanagedType.I4)]FileAttrib dwFlagsAndAttributes,
            [In]IntPtr hTemplateFile);


        //------------------------------------------------------------------------------------------
        // Reads data from a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool ReadFile([In]IntPtr hFile, [In]IntPtr lpBuffer,
            [In]uint nNumberOfBytesToRead, [Out]out uint lpNumberOfBytesRead,
            [In]IntPtr lpOverlapped);


        //------------------------------------------------------------------------------------------
        // Writes data to a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool WriteFile([In]IntPtr hFile, [In]IntPtr lpBuffer, 
            [In]uint nNumberOfBytesToWrite, [Out]out uint lpNumberOfBytesWritten,
            [In]IntPtr lpOverlapped);


        //------------------------------------------------------------------------------------------
        // Sets the end of a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool SetEndOfFile([In]IntPtr hFile);


        //------------------------------------------------------------------------------------------
        // Flush the file buffers
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static extern bool FlushFileBuffers([In]IntPtr hFile);


        //------------------------------------------------------------------------------------------
        // Closes a handle
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static extern bool CloseHandle([In]IntPtr hObject);

        //---------------------------------------------------------------------------------------------------
        // Formats an error code to a string
        //---------------------------------------------------------------------------------------------------
        [DllImport("Kernel32.dll", SetLastError = true)]
        internal static extern uint FormatMessage(uint dwFlags, IntPtr lpSource,
            uint dwMessageId, uint dwLanguageId, ref IntPtr lpBuffer,
            uint nSize, IntPtr pArguments);

        //---------------------------------------------------------------------------------------------------
        // Frees local memory
        //---------------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern IntPtr LocalFree(IntPtr hMem);

        //------------------------------------------------------------------------------------------
        // Sends a control command to the device
        //------------------------------------------------------------------------------------------
        [DllImport("Kernel32.dll", SetLastError = false, CharSet = CharSet.Auto)]
        internal static extern int DeviceIoControl([In] IntPtr hDevice,
            [In] IOCtrl IoControlCode, [In] IntPtr InBuffer, [In] uint nInBufferSize,
            [Out] IntPtr OutBuffer, [In] uint nOutBufferSize, [Out][In] ref uint pBytesReturned,
            [In] IntPtr Overlapped);
    }
}
